import datetime
import signal
import tempfile
import os
import subprocess
from absl import app, flags
import jinja2
from tornado.ioloop import IOLoop
from itertools import zip_longest
from coin.tool.strat_monitor.app.monitor_runner_slurm.monitor_runner_slurm import SBatchRunner
from coin.exchange.base.strategy_logger import parse_file_name
from coin.support.proto_log.logic.util import enumerate_strat_log_into_stratinfos
from concurrent.futures import ProcessPoolExecutor, as_completed


FLAGS = flags.FLAGS


def _grouper(n, iterable, fillvalue=None):
  args = [iter(iterable)] * n
  return zip_longest(fillvalue=fillvalue, *args)


def _get_every_1_strategies(n, strat_infos):
  strategies = []
  for strat_info in strat_infos:
    if strat_info.strategy_name not in strategies:
      strategies.append(strat_info.strategy_name)
  return strategies


def _get_every_n_strategies(n, strat_infos):
  # TODO(daniel):
  strategies = []
  for strat_info in strat_infos:
    if strat_info.strategy_name not in strategies:
      strategies.append(strat_info.strategy_name)
  strategy_list = []
  for n0, n1, n2, n3, n4, n5, n6, n7, n8, n9 in _grouper(n, strategies):
    tmp = [n0, n1, n2, n3, n4, n5, n6, n7, n8, n9]
    strats = [i for i in tmp if i]
    names = ','.join(strats)
    info = {
        'name': names,
        'mem-per-cpu': '%dM' % (0),  # TODO(daniel):
    }
    strategy_list.append(info)
  return strategy_list


def _get_yesterday():
  return (datetime.date.today() - datetime.timedelta(days=1)).strftime('%Y%m%d')


def generate_summary_report(tmp_dir, trading_date, log_type):
  cmd = f"./pyrunner python/coin/support/proto_log/app/strat_state_anaylyzer.py" \
        f" --generate_summary_report_only" \
        f" --tmp_dir={tmp_dir}" \
        f" --output_dir={FLAGS.output_dir}" \
        f" --trading_date={trading_date}"
  ret = subprocess.run(cmd, shell=True, encoding="utf-8")
  status = 'fail'
  if ret.returncode == 0:
    status = 'success'
  print("generate summary report:", status)


def _execute(script):
  print(script)
  ret = subprocess.run(script, shell=True)
  #if ret.returncode != 0:
  #  raise f"fail to process: {script}"
  print("done")


def main(argv):
  # check dir of strat proto log
  strat_proto_log_dir = '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log'
  if not os.path.exists(strat_proto_log_dir):
    print('ERROR: dir non-exist %s' % strat_proto_log_dir)
    return
  if not os.path.exists(FLAGS.output_dir):
    os.makedirs(FLAGS.output_dir)
  # check trading date
  trading_date = FLAGS.trading_date
  if trading_date is None:
    trading_date = _get_yesterday()
  # collect strategy name
  strat_infos = enumerate_strat_log_into_stratinfos(strat_proto_log_dir,
                                                    trading_date,
                                                    parse_file_name)
  strategy_list = _get_every_1_strategies(1, strat_infos)
  # debug
  #strategy_list = [
  #    "strategy-01.us-west-2 lm_agg2_kraken_spot_usd_x3",
  #    "strategy-02.ap-northeast-1c.huobi lm_agg2_huobi_spot_3",
  #]

  scripts = []
  for strategy in strategy_list:
    #  ./pyrunner python/coin/support/proto_log/app/strat_state_anaylyzer.py \
    #    --slurm_run \
    #    --trading_date=20220512 \
    #    --strategy_name=basis_smm2_gdax_spot_dydx_swap \
    #    --tmp_dir=/tmp \
    cmds = [
        f"./pyrunner python/coin/support/proto_log/app/strat_state_anaylyzer.py",
        f"--slurm_run",
        f"--output_dir={FLAGS.output_dir}",
        f"--trading_date={trading_date}",
        f"--tmp_dir={FLAGS.tmp_dir}",
        f"--strategy_name={strategy}",
    ]
    scripts.append(" ".join(cmds))

  if FLAGS.dry_run:
    # dry run
    for script in scripts:
      print(script)
    return

  for s0, s1, s2, s3, s4, s5, s6, s7, s8, s9 in _grouper(10, scripts):
    with ProcessPoolExecutor(10) as executor:
      futures = []
      if s0:
        futures.append(executor.submit(_execute, s0))
      if s1:
        futures.append(executor.submit(_execute, s1))
      if s2:
        futures.append(executor.submit(_execute, s2))
      if s3:
        futures.append(executor.submit(_execute, s3))
      if s4:
        futures.append(executor.submit(_execute, s4))
      if s5:
        futures.append(executor.submit(_execute, s5))
      if s6:
        futures.append(executor.submit(_execute, s6))
      if s7:
        futures.append(executor.submit(_execute, s7))
      if s8:
        futures.append(executor.submit(_execute, s8))
      if s9:
        futures.append(executor.submit(_execute, s9))
      for fut in as_completed(futures):
        _ = fut.result()

  # generate summary report
  generate_summary_report(FLAGS.tmp_dir, trading_date, FLAGS.log_type)


if __name__ == '__main__':
  flags.DEFINE_string('output_dir', '/tmp/og_anaylyzer_local/output', 'output_dir.')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('tmp_dir', '/tmp/og_anaylyzer_local', 'dir for keep tmp report')
  flags.DEFINE_bool('dry_run', False, 'print sbatch script only and not run sbatch')
  flags.DEFINE_string('log_type', None, '')
  app.run(main)
